package com.simon.study.algorithm.review;

import com.simon.study.algorithm.tree.AVLTree;
import com.simon.study.algorithm.tree.TNode.RBNode;
import java.util.Objects;
import java.util.Optional;
/**
 * <p>
 *
 * @author simon
 */
public class RBTree {
    RBNode      root;
    int         size;


    public RBNode find(int data){
        return (RBNode) AVLTree.find(root, data);
    }

    public void add(int data){
        RBNode node = insert(root, data);
        if(size == 0){ root = node; }
        size++;
        fixAfterAdd(node);
    }

    private void fixAfterAdd(RBNode node){
        while (node != root && node.parent().getColor() == RBNode.RED){
            RBNode uncle = node.uncle(), parent = node.parent(), grand = node.grandparent();
            if(uncle != null && uncle.getColor() == RBNode.RED){
                uncle.setColor(RBNode.BLACK);
                parent.setColor(RBNode.BLACK);
                grand.setColor(RBNode.RED);
                node = grand; continue;
            }

            if(parent == grand.getLeft()){
                if(node == parent.getRight()){
                    leftRotate(parent);
                    node = parent; parent = node.parent();
                }
                grand.setColor(RBNode.RED);
                parent.setColor(RBNode.BLACK);
                rightRotate(grand);
            }else{
                if(node == parent.getLeft()){
                    rightRotate(parent);
                    node = parent; parent = node.parent();
                }
                grand.setColor(RBNode.RED);
                parent.setColor(RBNode.BLACK);
                leftRotate(grand);
            }
        }
        root.setColor(RBNode.BLACK);
    }



    private RBNode leftRotate(RBNode spot){
        RBNode node = (RBNode) AVLTree.leftRotate(spot);
        adjustAfterRotate(spot, node);
        if(root == spot){ root = node; }
        return node;
    }

    private RBNode rightRotate(RBNode spot){
        RBNode node = (RBNode) AVLTree.rightLeftRotate(spot);
        adjustAfterRotate(spot, node);

        if(root == spot){ root = node; }
        return node;
    }


    public void delete(int data){
        RBNode node = find(data);

        if(Objects.isNull(node)){ return; }

        size--;
        deleteNode(node);
    }



    private void deleteNode(RBNode node){
        if(node.getLeft() != null && node.getRight() != null){
            RBNode s = successor(node);
            node.setData(s.getData());

            node     = s;
        }

        RBNode target = (RBNode) Optional.ofNullable(node.getLeft()).orElse(node.getRight());
        RBNode parent = node.parent();

        if (target != null) {
            target.setParent( parent );

            if (parent == null) {
                root = target;
            }
            else if (node == parent.getLeft()) {
                parent.setLeft(target);
            }else {
                parent.setRight(target);
            }

            node.setParent(null).setLeft(null).setRight(null);
            if (node.getColor() == RBNode.BLACK) {
                fixAfterDelete(target);
            }
        }else{
            if(node.getColor() == RBNode.BLACK){
                if (parent == null) { root = null; }
                else {
                    fixAfterDelete(node);
                }
            }
            node.setParent(null).setLeft(null).setRight(null);
        }
    }

    private void fixAfterDelete(RBNode node){
        while (node != root && node.getColor() == RBNode.BLACK) {
            RBNode sibling = node.sibling(), parent = node.parent();

            if (node == parent.getLeft()) {
                if (sibling.getColor() == RBNode.RED) {
                    sibling.setColor(RBNode.BLACK);
                    parent.setColor(RBNode.RED);
                    leftRotate(parent);
                    sibling = node.sibling();
                }

                if (sibling.getLeft() != null && node.getRight() != null) {
                    if (sibling.getRight() == null) {
                        sibling.setColor(RBNode.RED);
                        ((RBNode)sibling.getLeft()).setColor(RBNode.BLACK);
                        rightRotate(sibling);
                        sibling = node.sibling();
                    }

                    sibling.setColor(parent.getColor());
                    parent.setColor(RBNode.BLACK);
                    ((RBNode)sibling.getRight()).setColor(RBNode.BLACK);
                    leftRotate(parent);
                    break;
                }else{
                    sibling.setColor(RBNode.RED);
                    node = parent;
                }
            }else{
                if (sibling.getColor() == RBNode.RED) {
                    sibling.setColor(RBNode.BLACK);
                    parent.setColor(RBNode.RED);
                    rightRotate(parent);
                    sibling = node.sibling();
                }

                if (sibling.getLeft() != null && sibling.getRight() != null) {
                    if (sibling.getLeft() == null) {
                        sibling.setColor(RBNode.RED);
                        ((RBNode)sibling.getRight()).setColor(RBNode.BLACK);
                        leftRotate(sibling);
                        sibling = node.sibling();
                    }

                    sibling.setColor(parent.getColor());
                    parent.setColor(RBNode.BLACK);
                    ((RBNode)sibling.getLeft()).setColor(RBNode.BLACK);
                    rightRotate(parent);
                    break;
                }else{
                    sibling.setColor(RBNode.RED);
                    node = parent;
                }
            }
        }
        node.setColor(RBNode.BLACK);
    }


    private RBNode successor(RBNode node){
        if (node.getRight() != null) {
            RBNode s = (RBNode) node.getRight();

            while (s.getLeft() != null) {
                s = (RBNode) s.getLeft();
            }

            return s;
        }
        else{
            RBNode s = node.parent();
            RBNode c = node;
            while (s != null && c == s.getLeft()) {
                c = s;
                s = s.parent();
            }
            return s;
        }
    }


    public static void adjustAfterRotate(RBNode spot, RBNode node){
        node.setParent(spot.parent());
        spot.setParent(node);

        Optional.ofNullable(node.parent()).map(x -> x.getLeft() == spot ? x.setLeft(node) : x.setRight(node));
        Optional.ofNullable(spot.getLeft()).ifPresent(x -> x.setParent(spot));
        Optional.ofNullable(spot.getRight()).ifPresent(x -> x.setParent(spot));
    }

    public static RBNode insert(RBNode node, int data){
        if(node == null){ return new RBNode(data, RBNode.RED); }
        RBNode target = null;
        if(data > node.getData()){
            target = insert((RBNode) node.getRight(), data);
            if(node.getRight() == null){ node.setRight(target); target.setParent(node); }
        }else{
            target = insert((RBNode) node.getLeft(), data);
            if(node.getLeft()  == null){ node.setLeft(target);  target.setParent(node); }
        }
        return target;
    }
}
